
import torch
import torch.nn.functional as F
from captum.attr import Lime, LimeBase
from captum._utils.models.linear_model import SkLearnLinearRegression, SkLearnLasso
import pandas as pd
from model.predict import bert_lstm_predict
from modules.embed import gen_emb
def forward_func(text):
    output,pred=bert_lstm_predict(text)
    return output

def exp_embedding_cosine_distance(original_inp, perturbed_inp, _, **kwargs):
    original_emb = gen_emb(original_inp)
    perturbed_emb = gen_emb(perturbed_inp)
    distance = 1 - F.cosine_similarity(original_emb, perturbed_emb, dim=1)
    return torch.exp(-1 * (distance ** 2) / 2)

def interp_to_input(interp_sample, original_input, **kwargs):
    return original_input[interp_sample.bool()].view(original_input.size(0), -1)

def bernoulli_perturb(text, **kwargs):
    probs = torch.ones_like(text) * 0.5
    return torch.bernoulli(probs).long()

lasso_lime_base = LimeBase(
    forward_func,
    interpretable_model=SkLearnLasso(alpha=0.08),
    similarity_func=exp_embedding_cosine_distance,
    perturb_func=bernoulli_perturb,
    perturb_interpretable_space=True,
    from_interp_rep_transform=interp_to_input,
    to_interp_rep_transform=None
)

